import os
import numpy as np
import pandas as pd
import netCDF4 as nc
from scipy.spatial import cKDTree
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
from tqdm import tqdm
import logging

# --- 全局常量定义 (Global Constants) ---

# 1. 时空匹配部分常量 (Constants for Matching)
CLOUDSAT_CSV = "/mnt/datastore/liudddata/cloudsat_data/cloudsat_cbh_csv/cloudsat_cbh_20200506.csv"
FY_ROOT_DIR = "/mnt/datastore/liudddata/result/20200506_droupout"
COORD_FILE = "FY4A_coordinates.nc"
# 匹配结果CSV的输出路径 (The matched data will still be saved to this CSV for your records)
OUTPUT_CSV = "/mnt/datastore/liudddata/cloudsat_data/cloudsat_cbh_csv/cloudsat_matched_2020_0506_all_hours.csv"

# 2. 绘图部分常量 (Constants for Plotting)
# 所有折线图的保存目录 (Directory to save all the generated line chart images)
PLOT_SAVE_DIR = '/mnt/datastore/liudddata/cloudsat_data/cloudsat_cbh_csv/match_line_chart_2020_0506'

# --- 日志和绘图设置 (Logging and Plotting Settings) ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# 设置图片清晰度 (Set plot resolution)
plt.rcParams['figure.dpi'] = 300


# --- 函数定义部分 (Function Definitions) ---

# ==============================================================================
# 时空匹配相关函数 (Functions for Spatio-temporal Matching)
# 这部分函数来自您的第一个脚本，无需修改
# ==============================================================================

def read_fy_data(filename):
    """读取风云数据文件 (Read FY data file)"""
    with nc.Dataset(filename, 'r') as f:
        cth = f.variables['cth'][:]
        cbh = f.variables['predicted_mean'][:]
        clt = f.variables['clt'][:]  # --- 新增 ---
        cth = np.where(cth == "--", np.nan, cth)
        cbh = np.where(cbh == "--", np.nan, cbh)
        clt = np.where(clt == "--", np.nan, clt)  # --- 新增 ---
    return cth, cbh, clt


def read_coordinate_data():
    """读取并预处理风云卫星坐标数据 (Read and preprocess FY satellite coordinate data)"""
    with nc.Dataset(COORD_FILE, 'r') as f:
        lat = f.variables['lat'][:, :].T
        lon = f.variables['lon'][:, :].T
        lat[np.isnan(lat)] = -90.0
        lon[np.isnan(lon)] = 360.0
    return lat, lon


def spherical_to_cartesian(lon, lat):
    """球坐标转笛卡尔坐标 (Spherical to Cartesian coordinates)"""
    lon_rad = np.deg2rad(lon)
    lat_rad = np.deg2rad(90.0 - lat)
    x = np.sin(lat_rad) * np.cos(lon_rad)
    y = np.sin(lat_rad) * np.sin(lon_rad)
    z = np.cos(lat_rad)
    return np.column_stack([x, y, z])


def build_fy_kdtree():
    """预处理风云坐标并构建KD树 (Preprocess FY coordinates and build KD-tree)"""
    lat_fy, lon_fy = read_coordinate_data()
    valid_mask = (~np.isnan(lat_fy)) & (~np.isnan(lon_fy))
    coords_cart = spherical_to_cartesian(lon_fy[valid_mask], lat_fy[valid_mask])
    return cKDTree(coords_cart), valid_mask, lat_fy, lon_fy


def match_observations(tree, valid_mask, lat_fy, lon_fy, cth_data, cbh_data,clt_data,  # --- 新增clt_data参数 ---
                       cloudsat_lons, cloudsat_lats, cloudsat_times,
                       cloudsat_cth, cloudsat_cbh, radius_km=2.0):
    """执行空间匹配的核心函数 (Core function for spatial matching)"""
    matches = []
    earth_radius = 6371.0
    radius_rad = radius_km / earth_radius
    cs_points = spherical_to_cartesian(cloudsat_lons, cloudsat_lats)
    indices = tree.query_ball_point(cs_points, r=radius_rad)
    valid_indices = np.where(valid_mask)

    for i, (cs_lon, cs_lat, cs_time, cs_cth, cs_cbh) in enumerate(zip(
            cloudsat_lons, cloudsat_lats, cloudsat_times, cloudsat_cth, cloudsat_cbh)):
        if not indices[i]:
            continue
        fy_linear_idx = indices[i][0]
        fy_row, fy_col = valid_indices[0][fy_linear_idx], valid_indices[1][fy_linear_idx]
        fy_lat = lat_fy[fy_row, fy_col]
        fy_lon = lon_fy[fy_row, fy_col]
        fy_cth = cth_data[fy_row, fy_col]
        fy_cbh = cbh_data[fy_row, fy_col]
        fy_clt = clt_data[fy_row, fy_col]  # --- 新增：提取匹配点的云类型 ---
        matches.append({
            'cloudsat_time': cs_time,
            'cloudsat_lat': cs_lat,
            'cloudsat_lon': cs_lon,
            'cloudsat_cth': cs_cth,
            'cloudsat_cbh': cs_cbh,
            'fy_lat': fy_lat,
            'fy_lon': fy_lon,
            'fy_cth': fy_cth if not np.isnan(fy_cth) else None,
            'fy_cbh': fy_cbh if not np.isnan(fy_cbh) else None,
            'fy_clt': fy_clt if not np.isnan(fy_clt) else None
        })
    return matches


def process_daily_data(start_date, end_date):
    """主处理函数 - 进行全天24小时匹配，并返回结果DataFrame (Main processing function - performs 24-hour matching and returns the resulting DataFrame)"""
    fy_tree, valid_mask, lat_fy, lon_fy = build_fy_kdtree()
    cloudsat_df = pd.read_csv(CLOUDSAT_CSV, parse_dates=['time'])
    all_matches = []

    current_date = start_date
    while current_date <= end_date:
        for hour in range(24):
            current_time = current_date.replace(hour=hour, minute=0, second=0, microsecond=0)
            fy_filename = os.path.join(
                FY_ROOT_DIR,
                current_time.strftime("%Y%m%d%H") + "_predicted_2d_mc.nc"
            )
            if not os.path.exists(fy_filename):
                continue
            try:
                cth_data, cbh_data, clt_data = read_fy_data(fy_filename)
                start_window = current_time
                end_window = current_time + timedelta(minutes=15)
                time_mask = ((cloudsat_df['time'] >= start_window) & (cloudsat_df['time'] < end_window))
                hourly_data = cloudsat_df[time_mask].copy()
                if hourly_data.empty:
                    continue
                matches = match_observations(
                    fy_tree, valid_mask, lat_fy, lon_fy, cth_data, cbh_data,clt_data,
                    hourly_data['longitude'].values, hourly_data['latitude'].values,
                    hourly_data['time'].values, hourly_data['cloudsat_cth'].values,
                    hourly_data['cloudsat_cbh'].values
                )
                if matches:
                    all_matches.extend(matches)
                    print(f"Processed {current_time.strftime('%Y-%m-%d %H:00')} with {len(matches)} matches")
            except Exception as e:
                print(f"Error processing {current_time.strftime('%Y-%m-%d %H:00')}: {str(e)}")
        current_date += timedelta(days=1)

    result_df = pd.DataFrame(all_matches)
    if result_df.empty:
        print(
            "\n在指定的日期范围内未找到任何匹配的数据点。 (No matching data points were found in the specified date range.)")
        return None  # 返回None表示没有匹配结果

    print("\n正在过滤匹配结果... (Filtering matched results...)")
    valid_mask = (result_df['cloudsat_cbh'].notna() & result_df['fy_cbh'].notna())
    filtered_df = result_df[valid_mask].copy()

    print(f"\n数据过滤报告: (Data Filtering Report:)")
    print(f"总匹配记录 (Total matched records): {len(result_df)}")
    print(f"有效记录(双CBH非空) (Valid records with non-null CBH): {len(filtered_df)}")
    if len(result_df) > 0:
        print(f"过滤率 (Filter rate): {100 * (1 - len(filtered_df) / len(result_df)):.1f}%")

    # 按照您的要求，脚本不再依赖于这个中间文件，但我们仍然保存它以备将来分析
    filtered_df.to_csv(OUTPUT_CSV, index=False)
    print(f"\nMatched results saved to {OUTPUT_CSV}")

    # **重要变更**: 返回处理和过滤后的DataFrame，以便直接用于绘图
    return filtered_df


# ==============================================================================
# 新增的绘图函数 (NEW: Plotting function, from your second script)
# ==============================================================================

def plot_matched_data(matched_df, save_dir):
    """
    接收一个包含匹配数据的DataFrame，并按小时绘制折线图。
    (Takes a DataFrame with matched data and plots line charts grouped by hour.)
    """
    print(f"\n开始生成图像并保存到目录 (Starting to generate plots and save to directory): {save_dir}")

    # 确保保存图片的目录存在
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
        print(f"已创建目录 (Created directory): {save_dir}")

    # 确保时间列是datetime类型 (作为独立函数，最好再确认一次)
    if not pd.api.types.is_datetime64_any_dtype(matched_df['cloudsat_time']):
        matched_df['cloudsat_time'] = pd.to_datetime(matched_df['cloudsat_time'])

    # 按小时对数据进行分组
    grouped = matched_df.groupby(matched_df['cloudsat_time'].dt.floor('H'))

    if grouped.ngroups == 0:
        print("没有可用于绘图的数据分组。 (No data groups available for plotting.)")
        return

    # 遍历每个小时的数据组并绘图，使用tqdm显示进度条
    for hour, group in tqdm(grouped, desc="正在生成图像 (Generating plots)"):
        # 创建一个新的图形
        plt.figure(figsize=(12, 7))

        # 为了让线条更清晰，按纬度给数据点排序
        group = group.sort_values(by='cloudsat_lat')

        # 绘制风云数据（红色）
        plt.plot(group['fy_lat'], group['fy_cth'], 'r-', label='FY CTH', marker='.', markersize=3, linestyle='-')
        plt.plot(group['fy_lat'], group['fy_cbh'], 'r--', label='FY CBH', marker='.', markersize=3, linestyle='--')

        # 绘制 CloudSat 数据（蓝色）
        plt.plot(group['cloudsat_lat'], group['cloudsat_cth'], 'b-', label='CloudSat CTH', marker='x', markersize=3,
                 linestyle='-')
        plt.plot(group['cloudsat_lat'], group['cloudsat_cbh'], 'b--', label='CloudSat CBH', marker='x', markersize=3,
                 linestyle='--')

        # 设置标题和坐标轴标签
        plt.title(f'Cloud Base and Top Height vs. Latitude ({hour.strftime("%Y-%m-%d %H:00 UTC")})')
        plt.xlabel('Latitude')
        plt.ylabel('Height (m)')

        # 设置Y轴范围，可以根据您的数据大致范围调整
        plt.ylim(0, 20000)

        # 显示图例和网格
        plt.legend()
        plt.grid(True, linestyle='--', alpha=0.6)
        plt.tight_layout()

        # 生成图片文件名并保存
        filename = os.path.join(save_dir, f'match_{hour.strftime("%Y-%m-%d_%H%M")}.png')
        plt.savefig(filename)
        plt.close()  # 关闭当前图形以释放内存

    print(f"\n绘图完成！共生成 {len(grouped)} 张图像。 (Plotting complete! Generated {len(grouped)} images.)")


# --- 主执行部分 (Main Execution Block) ---

if __name__ == "__main__":
    # 1. 设置要处理的日期范围
    start_date = datetime(2020, 5, 1)
    end_date = datetime(2020, 6, 30)

    # 2. 执行时空匹配任务
    print("--- 步骤 1: 开始执行时空匹配 (Step 1: Starting spatio-temporal matching) ---")
    final_matched_df = process_daily_data(start_date, end_date)

    # 3. 如果匹配成功，则执行绘图任务
    #    检查返回的DataFrame是否有效 (不是None且不为空)
    if final_matched_df is not None and not final_matched_df.empty:
        print(
            "\n--- 步骤 2: 匹配成功，开始根据匹配结果生成图像 (Step 2: Match successful, starting plot generation) ---")
        plot_matched_data(final_matched_df, PLOT_SAVE_DIR)
        print("\n所有任务已成功完成！ (All tasks completed successfully!)")
    else:
        print(
            "\n时空匹配未找到任何有效数据，因此不执行绘图步骤。 (No valid data found from matching, skipping plotting step.)")